# MLP模型，用于预测距离
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
import torch.nn.functional as F

class MlpModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(320000, 256, dtype=torch.float32),
            # nn.ReLU(),
            nn.ELU(),
            nn.Linear(256, 64, dtype=torch.float32),
            # nn.ReLU(),
            nn.ELU(),
            nn.Linear(64, 16, dtype=torch.float32),
            # nn.ReLU(),
            nn.ELU(),
            nn.Linear(16, 1, dtype=torch.float32)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits